package org.example;

import com.alibaba.csp.sentinel.Entry;
import com.alibaba.csp.sentinel.SphU;
import com.alibaba.csp.sentinel.slots.block.BlockException;
import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.alibaba.csp.sentinel.util.TimeUtil;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class FlowQpsDemo {

    // 定义一个常量作为资源标识
    private static final String KEY = "abc";

    // 使用原子整型来记录通过、被阻塞和总请求的次数
    private static AtomicInteger pass = new AtomicInteger();
    private static AtomicInteger block = new AtomicInteger();
    private static AtomicInteger total = new AtomicInteger();

    // volatile变量，用来控制程序是否停止
    private static volatile boolean stop = false;

    // 模拟的线程数
    //private static final int threadCount = 32;
    private static final int threadCount = 1;

    // 测试时的秒数，控制统计的时间
    private static int seconds = 30;

    public static void main(String[] args) throws Exception {
        // 初始化流控规则
        initFlowQpsRule();

        // 启动定时任务，用于统计
        tick();

        // 启动模拟流量的生成
        simulateTraffic();

        // 输出流控开始的提示
        System.out.println("===== begin to do flow control");
        System.out.println("only 20 requests per second can pass");
    }

    // 初始化QPS流控规则
    private static void initFlowQpsRule() {
        // 创建一个FlowRule的列表
        List<FlowRule> rules = new ArrayList<FlowRule>();
        // 创建一个FlowRule对象
        FlowRule rule1 = new FlowRule();
        rule1.setResource(KEY);  // 设置流控资源为KEY
        rule1.setCount(20);  // 设置QPS为20
        rule1.setGrade(RuleConstant.FLOW_GRADE_QPS);  // 设置流控的类型为QPS
        rule1.setLimitApp("default");  // 设置流控规则适用的应用
        rules.add(rule1);

        // 加载流控规则
        FlowRuleManager.loadRules(rules);
    }

    // 模拟流量的生成
    private static void simulateTraffic() {
        // 启动多个线程模拟请求
        for (int i = 0; i < threadCount; i++) {
            Thread t = new Thread(new RunTask());
            t.setName("simulate-traffic-Task");
            t.start();
        }
    }

    // 启动定时器任务，用于统计流量情况
    private static void tick() {
        Thread timer = new Thread(new TimerTask());
        timer.setName("sentinel-timer-task");
        timer.start();
    }

    // 定时器任务，用于每秒统计一次流量情况
    static class TimerTask implements Runnable {

        @Override
        public void run() {
            long start = System.currentTimeMillis();
            System.out.println("begin to statistic!!!");

            long oldTotal = 0;
            long oldPass = 0;
            long oldBlock = 0;

            // 在未停止的情况下，每秒统计一次数据
            while (!stop) {
                try {
                    TimeUnit.SECONDS.sleep(1);  // 每隔1秒统计一次
                } catch (InterruptedException e) {
                    // 忽略异常
                }

                // 获取当前的总请求数、通过请求数和被阻塞请求数
                long globalTotal = total.get();
                long oneSecondTotal = globalTotal - oldTotal;
                oldTotal = globalTotal;

                long globalPass = pass.get();
                long oneSecondPass = globalPass - oldPass;
                oldPass = globalPass;

                long globalBlock = block.get();
                long oneSecondBlock = globalBlock - oldBlock;
                oldBlock = globalBlock;

                // 输出每秒的统计数据
                System.out.println(seconds + " send qps is: " + oneSecondTotal);
                System.out.println(TimeUtil.currentTimeMillis() + ", total:" + oneSecondTotal
                        + ", pass:" + oneSecondPass
                        + ", block:" + oneSecondBlock);

                // 如果统计时间到达0秒，则停止统计
                if (seconds-- <= 0) {
                    stop = true;
                }
            }

            // 输出统计过程的时间消耗
            long cost = System.currentTimeMillis() - start;
            System.out.println("time cost: " + cost + " ms");
            System.out.println("total:" + total.get() + ", pass:" + pass.get()
                    + ", block:" + block.get());
            System.exit(0);  // 结束程序
        }
    }

    // 模拟任务，用于模拟请求流量的产生
    static class RunTask implements Runnable {
        @Override
        public void run() {
            while (!stop) {
                Entry entry = null;

                try {
                    // 请求进入流控，返回entry表示请求成功
                    entry = SphU.entry(KEY);
                    // 如果成功获得token，表示请求通过
                    pass.addAndGet(1);
                } catch (BlockException e1) {
                    // 如果被流控阻塞，增加阻塞计数
                    block.incrementAndGet();
                } catch (Exception e2) {
                    // 处理业务异常
                } finally {
                    // 总请求数增加
                    total.incrementAndGet();
                    // 请求完成后释放entry
                    if (entry != null) {
                        entry.exit();
                    }
                }

                // 模拟每个请求的延时
                Random random2 = new Random();
                try {
                    TimeUnit.MILLISECONDS.sleep(random2.nextInt(50));  // 随机休眠0-50毫秒
                } catch (InterruptedException e) {
                    // 忽略异常
                }
            }
        }
    }
}
